import cv2
import numpy as np
import pyautogui
import win32gui
from ultralytics import YOLO

# 遍历所有窗口，查找包含指定关键字的窗口
def find_window_by_keyword(keyword):
    def callback(hwnd, hwnds):
        window_title = win32gui.GetWindowText(hwnd)
        if keyword in window_title:
            hwnds.append(hwnd)
        return True

    hwnds = []
    win32gui.EnumWindows(callback, hwnds)
    if hwnds:
        return hwnds[0]
    return None

# 获取窗口内的截图
def capture_window(hwnd):
    left, top, right, bottom = win32gui.GetWindowRect(hwnd)
    width = right - left
    height = bottom - top
    screenshot = pyautogui.screenshot(region=(left, top, width, height))
    screenshot = cv2.cvtColor(np.array(screenshot), cv2.COLOR_RGB2BGR)
    return screenshot

# 加载模型
model = YOLO("G:/yolov8/yolo/v11/runs/detect/train20/weights/best.pt")

# 查找窗口
window_keyword = "梦幻西游"
hwnd = find_window_by_keyword(window_keyword)
if hwnd:
    max_retries = 10  # 最大重试次数
    retry_count = 0
    while retry_count < max_retries:
        # 获取窗口截图
        screenshot = capture_window(hwnd)

        # 进行检测
        results = model.predict(source=screenshot, show=False, save=False, classes=[0, 9])

        # 假设只有一个检测结果
        result = results[0]
        boxes = result.boxes

        # 初始化变量
        x, y, c, d = None, None, None, None
        a, b = None, None
        e, f = None, None

        # 查找类别 9 的目标
        for box in boxes:
            if int(box.cls) == 9:
                # 获取右下坐标
                x, y = box.xyxy[0][2].item(), box.xyxy[0][3].item()
                # 获取中心坐标
                c, d = box.xywh[0][0].item(), box.xywh[0][1].item()
                break

        if x is not None and y is not None:
            # 获取系统鼠标当前坐标
            current_x, current_y = pyautogui.position()
            # 计算距离
            distance_x = x - current_x
            distance_y = y - current_y
            # 移动鼠标 90% 的距离
            target_x = current_x + distance_x * 0.9
            target_y = current_y + distance_y * 0.9
            pyautogui.moveTo(target_x, target_y, duration=0.12)
            m, n = pyautogui.position()
            print(f"鼠标已朝类别 9 右下坐标移动 90% 的距离，当前坐标: ({m}, {n})")

            # 查找类别 0 的目标
            for box in boxes:
                if int(box.cls) == 0:
                    # 获取左上角坐标
                    a, b = box.xyxy[0][0].item(), box.xyxy[0][1].item()
                    break

            if a is not None and b is not None:
                # 计算新坐标
                new_x = m + x - a
                new_y = n + y - b
                # 移动鼠标并点击
                pyautogui.moveTo(new_x, new_y, duration=0.1)
                pyautogui.click()
                print(f"鼠标已移动到新坐标 ({new_x}, {new_y}) 并点击一次")

                # 再次进行检测获取类别 0 的新坐标
                results = model.predict(source=screenshot, show=False, save=False, classes=[0])
                result = results[0]
                boxes = result.boxes
                for box in boxes:
                    if int(box.cls) == 0:
                        e, f = box.xyxy[0][0].item(), box.xyxy[0][1].item()
                        break

                if e is not None and f is not None:
                    # 计算最终坐标
                    final_x = m + x - a + c - e
                    final_y = n + y - b + d - f
                    # 移动鼠标并点击
                    pyautogui.moveTo(final_x, final_y, duration=0.1)
                    pyautogui.click()
                    print(f"鼠标已移动到新坐标 ({final_x}, {final_y}) 并点击一次")
            break
        else:
            print("未找到类别 9，按下 F9 刷新，继续检测...")
            pyautogui.press('f9')
            retry_count += 1
            pyautogui.sleep(2)  # 等待 2 秒后再次检测
    if retry_count == max_retries:
        print("达到最大重试次数，仍未找到类别 9。")
else:
    print(f"未找到包含关键字 '{window_keyword}' 的窗口")